import requests
import os
import argparse
from tqdm import tqdm

def get_download_links():
    """获取下载链接列表"""
    url = "https://net-secondary.web.minecraft-services.net/api/v1.0/download/links"
    
    headers = {
        "Host": "net-secondary.web.minecraft-services.net",
        "Connection": "keep-alive",
        "sec-ch-ua-platform": "\"Windows\"",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36",
        "sec-ch-ua": "\"Not)A;Brand\";v=\"8\", \"Chromium\";v=\"138\", \"Google Chrome\";v=\"138\"",
        "content-type": "application/json",
        "sec-ch-ua-mobile": "?0",
        "Accept": "*/*",
        "Origin": "https://www.minecraft.net",
        "Sec-Fetch-Site": "cross-site",
        "Sec-Fetch-Mode": "cors",
        "Sec-Fetch-Dest": "empty",
        "Referer": "https://www.minecraft.net/",
        "Accept-Encoding": "identity",
        "Accept-Language": "zh-CN,zh;q=0.9"
    }
    
    resp = requests.get(url, headers=headers)
    resp.raise_for_status()
    data = resp.json()
    
    return data["result"]["links"]

def get_filename_by_type(download_type):
    """根据下载类型返回标准文件名"""
    filename_map = {
        "serverBedrockWindows": "bedrock-server_windows.zip",
        "serverBedrockLinux": "bedrock-server_linux.zip",
        "serverBedrockPreviewWindows": "bedrock-server-preview_windows.zip",
        "serverBedrockPreviewLinux": "bedrock-server-preview_linux.zip"
    }
    return filename_map.get(download_type)

def format_size(size_bytes):
    """格式化文件大小"""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} TB"

def download_file(url, filename, show_progress=True, max_retries=3, download_dir=None):
    """下载文件到指定目录，支持断点续传和重试"""
    
    # 如果指定了下载目录，则拼接完整路径
    if download_dir:
        filepath = os.path.join(download_dir, filename)
    else:
        filepath = filename
    
    for retry in range(max_retries):
        try:
            # 检查是否已存在部分下载的文件
            resume_pos = 0
            if os.path.exists(filepath):
                resume_pos = os.path.getsize(filepath)
            
            # 设置请求头，支持断点续传
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }
            if resume_pos > 0:
                headers['Range'] = f'bytes={resume_pos}-'
            
            # 开始下载，设置更长的超时时间
            response = requests.get(url, headers=headers, stream=True, timeout=(30, 300))
            response.raise_for_status()
            
            # 获取文件总大小
            if response.status_code == 206:  # 部分内容（断点续传）
                total_size = resume_pos + int(response.headers.get('content-length', 0))
            else:
                total_size = int(response.headers.get('content-length', 0))
                resume_pos = 0  # 服务器不支持断点续传，从头开始
            
            # 显示下载信息
            if retry > 0:
                print(f"  重试 {retry}/{max_retries}...")
            
            if not show_progress:
                print(f"下载中: {filename} ({format_size(total_size)})")
            
            # 下载并保存文件
            mode = 'ab' if resume_pos > 0 else 'wb'
            with open(filepath, mode) as f:
                if show_progress and total_size > 0:
                    # 显示进度条
                    with tqdm(total=total_size, initial=resume_pos, unit='B', 
                             unit_scale=True, desc=filename, ncols=100) as pbar:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                                pbar.update(len(chunk))
                else:
                    # 不显示进度条
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            
            if not show_progress:
                print(f"下载完毕: {filename}")
            
            return True
            
        except (requests.exceptions.ConnectionError, 
                requests.exceptions.ChunkedEncodingError,
                requests.exceptions.Timeout) as e:
            if retry < max_retries - 1:
                print(f"  连接中断，准备重试... ({str(e)[:50]})")
                continue
            else:
                print(f"下载失败 {filename}: 已达最大重试次数")
                # 不删除文件，保留用于下次断点续传
                return False
                
        except Exception as e:
            print(f"下载失败 {filename}: {str(e)}")
            # 其他错误则删除不完整的文件
            if os.path.exists(filepath):
                os.remove(filepath)
            return False
    
    return False

def main():
    parser = argparse.ArgumentParser(description='Minecraft 基岩版服务端下载工具')
    parser.add_argument('--no-progress', action='store_true', 
                       help='不显示进度条')
    parser.add_argument('--types', nargs='+', 
                       choices=['windows', 'linux', 'preview-windows', 'preview-linux'],
                       help='指定要下载的版本类型（默认下载所有）')
    parser.add_argument('--output-dir', type=str, default=None,
                       help='指定下载目录（默认为当前目录）')
    
    args = parser.parse_args()
    show_progress = not args.no_progress
    
    # 确定下载目录
    download_dir = args.output_dir
    
    # 如果在 Docker 环境中且没有指定输出目录，检查是否存在 /app/downloads
    if not download_dir and os.path.exists('/app/downloads'):
        download_dir = '/app/downloads'
    
    # 创建下载目录（如果不存在）
    if download_dir and not os.path.exists(download_dir):
        os.makedirs(download_dir, exist_ok=True)
    
    # 显示下载目录信息
    if download_dir:
        print(f"下载目录: {download_dir}")
    
    try:
        # 获取下载链接
        print("正在获取下载链接...")
        links = get_download_links()
        
        # 筛选要下载的文件
        download_list = []
        for item in links:
            download_type = item['downloadType']
            filename = get_filename_by_type(download_type)
            
            # 只处理基岩版服务端，跳过其他类型（如 Java 版 serverJar）
            if not filename:
                continue
            
            # 根据用户指定的类型筛选
            if args.types:
                type_match = False
                if 'windows' in args.types and download_type == 'serverBedrockWindows':
                    type_match = True
                elif 'linux' in args.types and download_type == 'serverBedrockLinux':
                    type_match = True
                elif 'preview-windows' in args.types and download_type == 'serverBedrockPreviewWindows':
                    type_match = True
                elif 'preview-linux' in args.types and download_type == 'serverBedrockPreviewLinux':
                    type_match = True
                
                if not type_match:
                    continue
            
            download_list.append({
                'url': item['downloadUrl'],
                'filename': filename,
                'type': download_type
            })
        
        if not download_list:
            print("没有可下载的文件")
            return
        
        # 下载文件
        print(f"找到 {len(download_list)} 个基岩版服务端文件")
        print(f"开始下载...\n")
        success_count = 0
        
        for item in download_list:
            print(f"正在下载: {item['filename']} ({item['type']})")
            if download_file(item['url'], item['filename'], show_progress, download_dir=download_dir):
                success_count += 1
            print()
        
        # 显示下载结果
        print(f"\n下载完成！成功: {success_count}/{len(download_list)}")
        
    except Exception as e:
        print(f"错误: {str(e)}")
        return 1
    
    return 0

if __name__ == "__main__":
    exit(main())
